# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#

import logging
import os
import time
import json
import argparse
from collections import OrderedDict
import numpy as np
import torch

from src.utils import bool_flag, initialize_exp
from src.models import build_model
from src.trainer import Trainer
from src.evaluation import Evaluator


VALIDATION_METRIC = 'mean_cosine-csls_knn_10-S2T-10000'
DROPDICT_METRIC = 'mean_cosine-csls_knn_10-S2T-10000'


# main
parser = argparse.ArgumentParser(description='Unsupervised training')
parser.add_argument("--seed", type=int, default=-1, help="Initialization seed")
parser.add_argument("--verbose", type=int, default=2, help="Verbose level (2:debug, 1:info, 0:warning)")
parser.add_argument("--exp_path", type=str, default="", help="Where to store experiment logs and models")
parser.add_argument("--exp_name", type=str, default="debug", help="Experiment name")
parser.add_argument("--exp_id", type=str, default="", help="Experiment ID")
parser.add_argument("--cuda",  help="Whether to run on GPU", action="store_true")
parser.add_argument("--export", type=str, default="txt", help="Export embeddings after training (txt / pth)")
# data
parser.add_argument("--src_lang", type=str, default='en', help="Source language")
parser.add_argument("--tgt_lang", type=str, default='es', help="Target language")
parser.add_argument("--emb_dim", type=int, default=300, help="Embedding dimension")
parser.add_argument("--max_vocab", type=int, default=200000, help="Maximum vocabulary size (-1 to disable)")
# mapping
parser.add_argument("--map_id_init", type=bool_flag, default=True, help="Initialize the mapping as an identity matrix")
parser.add_argument("--map_beta", type=float, default=0.001, help="Beta for orthogonalization")
parser.add_argument("--map_bias_decay", type=float, default=0.9, help="bias decay for mapping.")
parser.add_argument("--map_bias", type=bool_flag, default=False, help="Used mapping bias to instruct orthogonalization")
# discriminator
parser.add_argument("--dis_layers", type=int, default=2, help="Discriminator layers")
parser.add_argument("--dis_hid_dim", type=int, default=2048, help="Discriminator hidden layer dimensions")
parser.add_argument("--dis_dropout", type=float, default=0., help="Discriminator dropout")
parser.add_argument("--dis_input_dropout", type=float, default=0.1, help="Discriminator input dropout")
parser.add_argument("--dis_steps", type=int, default=5, help="Discriminator steps")
parser.add_argument("--dis_lambda", type=float, default=1, help="Discriminator loss feedback coefficient")
parser.add_argument("--dis_most_frequent", type=int, default=75000, help="Select embeddings of the k most frequent words for discrimination (0 to disable)")
parser.add_argument("--dis_smooth", type=float, default=0.1, help="Discriminator smooth predictions")
parser.add_argument("--dis_clip_weights", type=float, default=0, help="Clip discriminator weights (0 to disable)")
# training adversarial
parser.add_argument("--n_adversarials", type=int, default=1, help="Turns of random-starts adversarial training")
parser.add_argument("--n_epochs", type=int, default=5, help="Number of epochs")
parser.add_argument("--epoch_size", type=int, default=1000000, help="Iterations per epoch")
parser.add_argument("--batch_size", type=int, default=32, help="Batch size")
parser.add_argument("--map_optimizer", type=str, default="sgd,lr=0.1", help="Mapping optimizer")
parser.add_argument("--dis_optimizer", type=str, default="sgd,lr=0.1", help="Discriminator optimizer")
parser.add_argument("--lr_decay", type=float, default=0.98, help="Learning rate decay (SGD only)")
parser.add_argument("--min_lr", type=float, default=1e-6, help="Minimum learning rate (SGD only)")
parser.add_argument("--lr_shrink", type=float, default=0.5, help="Shrink the learning rate if the validation metric decreases (1 to disable)")
# training refinement
parser.add_argument("--n_refinement", type=int, default=5, help="Number of refinement iterations (0 to disable the refinement procedure)")
parser.add_argument("--refine_vocab_cutoff", type=int, default=-1, help="The vocabulary cutoff using for training the mapper. (-1 to disable).")
# dictionary creation parameters (for refinement)
parser.add_argument("--dico_eval", type=str, default="default", help="Path to evaluation dictionary")
parser.add_argument("--dico_method", type=str, default='csls_knn_10', help="Method used for dictionary generation (nn/invsm_beta_30/csls_knn_10)")
parser.add_argument("--dico_build", type=str, default='S2T', help="S2T,T2S,S2T|T2S,S2T&T2S")
parser.add_argument("--dico_threshold", type=float, default=1e-6, help="Threshold confidence for dictionary generation")
parser.add_argument("--dico_max_rank", type=int, default=15000, help="Maximum dictionary words rank (0 to disable)")
parser.add_argument("--dico_min_size", type=int, default=0, help="Minimum generated dictionary size (0 to disable)")
parser.add_argument("--dico_max_size", type=int, default=0, help="Maximum generated dictionary size (0 to disable)")
# reload pre-trained embeddings
parser.add_argument("--src_emb", type=str, default="", help="Reload source embeddings")
parser.add_argument("--tgt_emb", type=str, default="", help="Reload target embeddings")
parser.add_argument("--normalize_embeddings", type=str, default="", help="Normalize embeddings before training")
# dropout prob decreasing by:
parser.add_argument("--use_sdi", help="whether to use stochastic dictionary induction", action="store_true")
parser.add_argument("--seeddict_sdi", help="whether to build a small seed dictionary to start.", action="store_true")
parser.add_argument("--sdi_cutoff", type=int, default=-1, help="Restrict the most frequent word to stochasticly drop off.")
# parser.add_argument("--dico_dropout_deps", type=str, default="precision_at_1-csls_knn_10", help="Stochastic dictionary dropout when deps-metrics is stable.")
parser.add_argument("--initial_keep_prob", type=float, default=0.1, help="The initial keep probility of dropout in stochastic dictionary induction.")
parser.add_argument("--stochastic_interval", type=int, default=10, help="The number of iterations that a prob to keep at least.")
parser.add_argument("--stochastic_multiplier", type=int, default=2.0 , help="stochastic dictionary induction multiplier (defaults to 2.0)")
parser.add_argument("--load_temp_model", type=str, default="", help="interrupt training last time and restart from checkpoint.")
# interaction with others
parser.add_argument("--return_GAN_xwzw", help="Return xw=mappping(x), zw=z after GAN initialization.", action="store_true")

def main():
    # parse parameters
    params = parser.parse_args()

    # check parameters
    assert not params.cuda or torch.cuda.is_available()
    assert 0 <= params.dis_dropout < 1
    assert 0 <= params.dis_input_dropout < 1
    assert 0 <= params.dis_smooth < 0.5
    assert params.dis_lambda > 0 and params.dis_steps > 0
    assert 0 < params.lr_shrink <= 1
    assert os.path.isfile(params.src_emb)
    assert os.path.isfile(params.tgt_emb)
    assert params.dico_eval == 'default' or os.path.isfile(params.dico_eval)
    assert params.export in ["", "txt", "pth"]

    # logging level
    logging.basicConfig(level=logging.DEBUG)
    # build model / trainer / evaluator
    logger = initialize_exp(params)
    src_emb, tgt_emb, mapping, discriminator = build_model(params, True)
    trainer = Trainer(src_emb, tgt_emb, mapping, discriminator, params)
    evaluator = Evaluator(trainer)

    if params.load_temp_model:
        trainer.reload_best(params.load_temp_model)
        to_log = OrderedDict({'n_epoch': -1})
        evaluator.dist_mean_cosine(to_log)
        # evaluator.all_eval(to_log)
        # JSON log / save best model / end of epoch
        logger.debug("__log__:%s" % json.dumps(to_log))

    """
    Learning loop for Adversarial Training
    """
    if params.n_adversarials:
        
        for adversarial in range(params.n_adversarials):
            logger.debug(f'----> ADVERSARIAL TRAINING {adversarial}<----\n')

            trainer.reset_models()
            # training loop
            for n_epoch in range(params.n_epochs):

                logger.debug('Starting adversarial training epoch %i...' % n_epoch)
                tic = time.time()
                n_words_proc = 0
                stats = {
                    'DIS_COSTS': [],
                    'MAP_BIAS_NORM': [],
                }

                for n_iter in range(0, params.epoch_size, params.batch_size):

                    # discriminator training
                    for _ in range(params.dis_steps):
                        trainer.dis_step(stats)

                    # mapping training (discriminator fooling)
                    n_words_proc += trainer.mapping_step(stats)

                    # log stats
                    if n_iter % 60000 == 0:
                        stats_str = [('DIS_COSTS', 'Discriminator loss'), ('MAP_BIAS_NORM','map bias norm')]
                        stats_log = ['%s: %.4f' % (v, np.mean(stats[k]))
                                    for k, v in stats_str if len(stats[k]) > 0]
                        stats_log.append('%i samples/s' % int(n_words_proc / (time.time() - tic)))
                        logger.debug(('%06i - ' % n_iter) + ' - '.join(stats_log))

                        # reset
                        tic = time.time()
                        n_words_proc = 0
                        for k, _ in stats_str:
                            del stats[k][:]

                # embeddings / discriminator evaluation
                to_log = OrderedDict({'n_epoch': n_epoch})
                # evaluator.all_eval(to_log)
                # evaluator.eval_dis(to_log)
                evaluator.dist_mean_cosine(to_log)


                # JSON log / save best model / end of epoch
                logger.debug("__log__:%s" % json.dumps(to_log))
                trainer.save_best(to_log, VALIDATION_METRIC)
                logger.debug('End of epoch %i.\n\n' % n_epoch)

                # update the learning rate (stop if too small)
                trainer.update_lr(to_log, VALIDATION_METRIC)
                if trainer.map_optimizer.param_groups[0]['lr'] < params.min_lr:
                    logger.debug('Learning rate < 1e-6. BREAK.')
                    break

            evaluator.all_eval(to_log)
            evaluator.eval_dis(to_log)
        
        if params.return_GAN_xwzw:
            trainer.build_dictionary()
            src_dico = trainer.src_dico
            tgt_dico = trainer.tgt_dico

            x = trainer.mapping(trainer.src_emb.weight).data
            z = trainer.tgt_emb.weight.data
            return (list(src_dico.word2id.keys()), src_dico.word2id, trainer.dico[:, 0], x,\
                    list(tgt_dico.word2id.keys()), tgt_dico.word2id, trainer.dico[:, 1], z )


    """
    Learning loop for Procrustes Iterative Refinement
    """
    if params.use_sdi:
        # Get the best mapping according to DROPDICT_METRIC
        logger.debug('----> ITERATIVE PROCRUSTES REFINEMENT WITH STOCHASTIC DICTIONARY INDUCTION<----\n\n')
        # if params.load_temp_model:
        #     trainer.reload_best()

        if params.seeddict_sdi:
            trainer.build_seed_dictionary()
            trainer.procrustes()
        
        # return trainer.mapping(trainer.src_emb.weight.data[:params.sdi_cutoff]), trainer.tgt_emb.weight.data[:params.sdi_cutoff]
        n_iter = 0
        keep_prob = params.initial_keep_prob
        last_improvement = 0
        best_objective = float('-inf')
        while True:
            logger.debug('Starting refinement iteration %i...' % n_iter)
                
            # build a dictionary from aligned embeddings
            trainer.build_dictionary(dropout_prob=1 - keep_prob, cutoff=params.sdi_cutoff)

            # apply the Procrustes solution
            trainer.procrustes()

            # embeddings evaluation
            to_log = OrderedDict({'n_iter': n_iter})
            # evaluator.all_eval(to_log)
            # evaluator.word_translation(to_log)
            evaluator.dist_mean_cosine(to_log, methods = ['csls_knn_10'], cutoff=params.sdi_cutoff)

            # JSON log / save best model / end of epoch
            # logger.debug("__log__:%s" % json.dumps(to_log))
            trainer.save_best(to_log, VALIDATION_METRIC)
            # logger.debug('End of refinement iteration %i.\n\n' % n_iter)

            # To decrease dropout-prob and check to break
            n_iter += 1
            objective = to_log[DROPDICT_METRIC]
            if objective - params.dico_threshold > best_objective:
                last_improvement = n_iter
                best_objective = objective
                logger.debug(f"__log__:{DROPDICT_METRIC} steps to {objective}")
            if n_iter - last_improvement > params.stochastic_interval:
                trainer.build_dictionary()
                evaluator.all_eval(to_log)
                logger.debug("__log__:%s" % json.dumps(to_log))
                if keep_prob >= 1.0:
                    break
                keep_prob = min(1.0, params.stochastic_multiplier*keep_prob)
                last_improvement = n_iter
                logging.info(f"__log__:keep_prob change to {keep_prob}")
    elif params.n_refinement > 0:
        # Get the best mapping according to VALIDATION_METRIC
        logger.debug('----> ITERATIVE PROCRUSTES REFINEMENT <----\n\n')
        if not params.load_temp_model:
            trainer.reload_best()

        # training loop
        # refine_cutoffs = [params.dis_most_frequent] * params.n_refinement +\
        #     list(range(params.dis_most_frequent, params.max_vocab+1, 10000))
        refine_cutoffs = [params.refine_vocab_cutoff] * params.n_refinement

        # for n_iter in range(params.n_refinement):
        for n_iter, refine_cutoff in enumerate(refine_cutoffs):

            logger.debug('Starting refinement iteration %i...' % n_iter)

            # build a dictionary from aligned embeddings
            trainer.build_dictionary(cutoff=refine_cutoff)

            # apply the Procrustes solution
            trainer.procrustes()

            # embeddings evaluation
            to_log = OrderedDict({'n_iter': n_iter, 'refine_cutoff': refine_cutoff})
            evaluator.all_eval(to_log)

            # JSON log / save best model / end of epoch
            logger.debug("__log__:%s" % json.dumps(to_log))
            trainer.save_best(to_log, VALIDATION_METRIC)
            logger.debug('End of refinement iteration %i.\n\n' % n_iter)

    # export embeddings
    if params.export:
        trainer.reload_best()
        trainer.export()

if __name__ == '__main__':
    main()